// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/record_batch.h>

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <memory>
#include <sstream>
#include <string>
#include <utility>

#include <nebula/core/array.h>
#include <nebula/array/concatenate.h>
#include <nebula/array/validate.h>
#include <nebula/core/pretty_print.h>
#include <turbo/utility/status.h>
#include <nebula/core/table.h>
#include <nebula/types/tensor.h>
#include <nebula/types/type.h>
#include <turbo/functional/iterator.h>
#include <turbo/log/logging.h>
#include <nebula/util/vector.h>
#include <nebula/core/visit_type_inline.h>

namespace nebula {

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::add_column(
            int i, std::string field_name, const std::shared_ptr<Array> &column) const {
        auto field = ::nebula::field(std::move(field_name), column->type());
        return add_column(i, field, column);
    }

    std::shared_ptr<Array> RecordBatch::get_column_by_name(const std::string &name) const {
        auto i = schema_->get_field_index(name);
        return i == -1 ? nullptr : column(i);
    }

    int RecordBatch::num_columns() const { return schema_->num_fields(); }

    /// \class SimpleRecordBatch
    /// \brief A basic, non-lazy in-memory record batch
    class SimpleRecordBatch : public RecordBatch {
    public:
        SimpleRecordBatch(std::shared_ptr<Schema> schema, int64_t num_rows,
                          std::vector<std::shared_ptr<Array>> columns,
                          std::shared_ptr<Device::SyncEvent> sync_event = nullptr)
                : RecordBatch(std::move(schema), num_rows),
                  boxed_columns_(std::move(columns)),
                  device_type_(DeviceAllocationType::kCPU),
                  sync_event_(std::move(sync_event)) {
            if (boxed_columns_.size() > 0) {
                device_type_ = boxed_columns_[0]->device_type();
            }

            columns_.resize(boxed_columns_.size());
            for (size_t i = 0; i < columns_.size(); ++i) {
                columns_[i] = boxed_columns_[i]->data();
                DKCHECK(device_type_ == columns_[i]->device_type());
            }
        }

        SimpleRecordBatch(const std::shared_ptr<Schema> &schema, int64_t num_rows,
                          std::vector<std::shared_ptr<ArrayData>> columns,
                          DeviceAllocationType device_type = DeviceAllocationType::kCPU,
                          std::shared_ptr<Device::SyncEvent> sync_event = nullptr)
                : RecordBatch(std::move(schema), num_rows),
                  columns_(std::move(columns)),
                  device_type_(device_type),
                  sync_event_(std::move(sync_event)) {
            boxed_columns_.resize(schema_->num_fields());
        }

        const std::vector<std::shared_ptr<Array>> &columns() const override {
            for (int i = 0; i < num_columns(); ++i) {
                // Force all columns to be boxed
                column(i);
            }
            return boxed_columns_;
        }

        std::shared_ptr<Array> column(int i) const override {
            std::shared_ptr<Array> result = std::atomic_load(&boxed_columns_[i]);
            if (!result) {
                result = make_array(columns_[i]);
                std::atomic_store(&boxed_columns_[i], result);
            }
            return result;
        }

        std::shared_ptr<ArrayData> column_data(int i) const override { return columns_[i]; }

        const ArrayDataVector &column_data() const override { return columns_; }

        turbo::Result<std::shared_ptr<RecordBatch>> add_column(
                int i, const std::shared_ptr<Field> &field,
                const std::shared_ptr<Array> &column) const override {
            KCHECK(field != nullptr);
            KCHECK(column != nullptr);
            KCHECK(column->device_type() == device_type_);

            if (!field->type()->equals(column->type())) {
                return turbo::failed_precondition_error("Column data type ", field->type()->name(),
                                                        " does not match field data type ",
                                                        column->type()->name());
            }
            if (column->length() != num_rows_) {
                return turbo::invalid_argument_error(
                        "Added column's length must match record batch's length. Expected length ",
                        num_rows_, " but got length ", column->length());
            }

            TURBO_MOVE_OR_RAISE(auto new_schema, schema_->add_field(i, field));
            return RecordBatch::create(std::move(new_schema), num_rows_,
                                       internal::AddVectorElement(columns_, i, column->data()),
                                       device_type_, sync_event_);
        }

        turbo::Result<std::shared_ptr<RecordBatch>> set_column(
                int i, const std::shared_ptr<Field> &field,
                const std::shared_ptr<Array> &column) const override {
            KCHECK(field != nullptr);
            KCHECK(column != nullptr);
            KCHECK(column->device_type() == device_type_);

            if (!field->type()->equals(column->type())) {
                return turbo::failed_precondition_error("Column data type ", field->type()->name(),
                                                        " does not match field data type ",
                                                        column->type()->name());
            }
            if (column->length() != num_rows_) {
                return turbo::invalid_argument_error(
                        "Added column's length must match record batch's length. Expected length ",
                        num_rows_, " but got length ", column->length());
            }

            TURBO_MOVE_OR_RAISE(auto new_schema, schema_->set_field(i, field));
            return RecordBatch::create(std::move(new_schema), num_rows_,
                                       internal::ReplaceVectorElement(columns_, i, column->data()),
                                       device_type_, sync_event_);
        }

        turbo::Result<std::shared_ptr<RecordBatch>> remove_column(int i) const override {
            TURBO_MOVE_OR_RAISE(auto new_schema, schema_->remove_field(i));
            return RecordBatch::create(std::move(new_schema), num_rows_,
                                       internal::DeleteVectorElement(columns_, i), device_type_,
                                       sync_event_);
        }

        std::shared_ptr<RecordBatch> replace_schema_metadata(
                const std::shared_ptr<const KeyValueMetadata> &metadata) const override {
            auto new_schema = schema_->with_metadata(metadata);
            return RecordBatch::create(std::move(new_schema), num_rows_, columns_, device_type_,
                                       sync_event_);
        }

        std::shared_ptr<RecordBatch> slice(int64_t offset, int64_t length) const override {
            std::vector<std::shared_ptr<ArrayData>> arrays;
            arrays.reserve(num_columns());
            for (const auto &field: columns_) {
                arrays.emplace_back(field->slice(offset, length));
            }
            int64_t num_rows = std::min(num_rows_ - offset, length);
            return std::make_shared<SimpleRecordBatch>(schema_, num_rows, std::move(arrays),
                                                       device_type_, sync_event_);
        }

        turbo::Status validate() const override {
            if (static_cast<int>(columns_.size()) != schema_->num_fields()) {
                return turbo::invalid_argument_error("Number of columns did not match schema");
            }
            return RecordBatch::validate();
        }

        const std::shared_ptr<Device::SyncEvent> &GetSyncEvent() const override {
            return sync_event_;
        }

        DeviceAllocationType device_type() const override { return device_type_; }

    private:
        std::vector<std::shared_ptr<ArrayData>> columns_;

        // Caching boxed array data
        mutable std::vector<std::shared_ptr<Array>> boxed_columns_;

        // the type of device that the buffers for columns are allocated on.
        // all columns should be on the same type of device.
        DeviceAllocationType device_type_;
        std::shared_ptr<Device::SyncEvent> sync_event_;
    };

    RecordBatch::RecordBatch(const std::shared_ptr<Schema> &schema, int64_t num_rows)
            : schema_(schema), num_rows_(num_rows) {}

    std::shared_ptr<RecordBatch> RecordBatch::create(
            std::shared_ptr<Schema> schema, int64_t num_rows,
            std::vector<std::shared_ptr<Array>> columns,
            std::shared_ptr<Device::SyncEvent> sync_event) {
        DKCHECK_EQ(schema->num_fields(), static_cast<int>(columns.size()));
        return std::make_shared<SimpleRecordBatch>(std::move(schema), num_rows,
                                                   std::move(columns), std::move(sync_event));
    }

    std::shared_ptr<RecordBatch> RecordBatch::create(
            std::shared_ptr<Schema> schema, int64_t num_rows,
            std::vector<std::shared_ptr<ArrayData>> columns, DeviceAllocationType device_type,
            std::shared_ptr<Device::SyncEvent> sync_event) {
        DKCHECK_EQ(schema->num_fields(), static_cast<int>(columns.size()));
        return std::make_shared<SimpleRecordBatch>(std::move(schema), num_rows,
                                                   std::move(columns), device_type,
                                                   std::move(sync_event));
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::make_empty(
            std::shared_ptr<Schema> schema, MemoryPool *memory_pool) {
        ArrayVector empty_batch(schema->num_fields());
        for (int i = 0; i < schema->num_fields(); i++) {
            TURBO_MOVE_OR_RAISE(empty_batch[i],
                                MakeEmptyArray(schema->field(i)->type(), memory_pool));
        }
        return RecordBatch::create(std::move(schema), 0, std::move(empty_batch));
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::FromStructArray(
            const std::shared_ptr<Array> &array, MemoryPool *memory_pool) {
        if (array->type_id() != Type::STRUCT) {
            return turbo::failed_precondition_error("Cannot construct record batch from array of type ",
                                                    *array->type());
        }
        if (array->null_count() != 0 || array->offset() != 0) {
            // If the struct array has a validity map or offset we need to push those into
            // the child arrays via Flatten since the RecordBatch doesn't have validity/offset
            const std::shared_ptr<StructArray> &struct_array =
                    turbo::checked_pointer_cast<StructArray>(array);
            TURBO_MOVE_OR_RAISE(std::vector<std::shared_ptr<Array>> fields,
                                struct_array->Flatten(memory_pool));
            return create(nebula::schema(array->type()->fields()), array->length(),
                          std::move(fields));
        }
        return create(nebula::schema(array->type()->fields()), array->length(),
                      array->data()->child_data);
    }

    namespace {

        turbo::Status ValidateColumnLength(const RecordBatch &batch, int i) {
            const auto &array = *batch.column(i);
            if (TURBO_UNLIKELY(array.length() != batch.num_rows())) {
                return turbo::invalid_argument_error("Number of rows in column ", i,
                                                     " did not match batch: ", array.length(), " vs ",
                                                     batch.num_rows());
            }
            return turbo::OkStatus();
        }

    }  // namespace

    turbo::Result<std::shared_ptr<StructArray>> RecordBatch::ToStructArray() const {
        if (num_columns() != 0) {
            // Only check the first column because `StructArray::create` already checks that the
            // child lengths are equal.
            TURBO_RETURN_NOT_OK(ValidateColumnLength(*this, 0));
            return StructArray::create(columns(), schema()->fields());
        }
        return std::make_shared<StructArray>(nebula::STRUCT({}), num_rows_,
                                             std::vector<std::shared_ptr<Array>>{},
                /*null_bitmap=*/nullptr,
                /*null_count=*/0,
                /*offset=*/0);
    }

    turbo::Result<std::shared_ptr<Tensor>> RecordBatch::ToTensor(bool null_to_nan, bool row_major,
                                                                 MemoryPool *pool) const {
        std::shared_ptr<Tensor> tensor;
        TURBO_RETURN_NOT_OK(
                internal::RecordBatchToTensor(*this, null_to_nan, row_major, pool, &tensor));
        return tensor;
    }

    const std::string &RecordBatch::column_name(int i) const {
        return schema_->field(i)->name();
    }

    bool RecordBatch::equals(const RecordBatch &other, bool check_metadata,
                             const EqualOptions &opts) const {
        if (num_columns() != other.num_columns() || num_rows_ != other.num_rows()) {
            return false;
        }

        if (!schema_->equals(*other.schema(), check_metadata)) {
            return false;
        }

        if (device_type() != other.device_type()) {
            return false;
        }

        for (int i = 0; i < num_columns(); ++i) {
            if (!column(i)->equals(other.column(i), opts)) {
                return false;
            }
        }

        return true;
    }

    bool RecordBatch::approx_equals(const RecordBatch &other, const EqualOptions &opts) const {
        if (num_columns() != other.num_columns() || num_rows_ != other.num_rows()) {
            return false;
        }

        if (device_type() != other.device_type()) {
            return false;
        }

        for (int i = 0; i < num_columns(); ++i) {
            if (!column(i)->approx_equals(other.column(i), opts)) {
                return false;
            }
        }

        return true;
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::replace_schema(
            std::shared_ptr<Schema> schema) const {
        if (schema_->num_fields() != schema->num_fields())
            return turbo::invalid_argument_error("RecordBatch schema fields", schema_->num_fields(),
                                                 ", did not match new schema fields: ", schema->num_fields());
        auto fields = schema_->fields();
        int n_fields = static_cast<int>(fields.size());
        for (int i = 0; i < n_fields; i++) {
            auto old_type = fields[i]->type();
            auto replace_type = schema->field(i)->type();
            if (!old_type->equals(replace_type)) {
                return turbo::invalid_argument_error(
                        "RecordBatch schema field index ", i, " type is ", old_type->to_string(),
                        ", did not match new schema field type: ", replace_type->to_string());
            }
        }
        return RecordBatch::create(std::move(schema), num_rows(), columns(), GetSyncEvent());
    }

    std::vector<std::string> RecordBatch::column_names() const {
        std::vector<std::string> names(num_columns());
        for (int i = 0; i < num_columns(); ++i) {
            names[i] = schema()->field(i)->name();
        }
        return names;
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::rename_columns(
            const std::vector<std::string> &names) const {
        int n = num_columns();

        if (static_cast<int>(names.size()) != n) {
            return turbo::invalid_argument_error("tried to rename a record batch of ", n, " columns but only ",
                                                 names.size(), " names were provided");
        }

        ArrayVector columns(n);
        FieldVector fields(n);

        for (int i = 0; i < n; ++i) {
            columns[i] = column(i);
            fields[i] = schema()->field(i)->with_name(names[i]);
        }

        return RecordBatch::create(::nebula::schema(std::move(fields)), num_rows(),
                                   std::move(columns), GetSyncEvent());
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::select_columns(
            const std::vector<int> &indices) const {
        int n = static_cast<int>(indices.size());

        FieldVector fields(n);
        ArrayVector columns(n);

        for (int i = 0; i < n; i++) {
            int pos = indices[i];
            if (pos < 0 || pos > num_columns() - 1) {
                return turbo::invalid_argument_error("Invalid column index ", pos, " to select columns.");
            }
            fields[i] = schema()->field(pos);
            columns[i] = column(pos);
        }

        auto new_schema =
                std::make_shared<nebula::Schema>(std::move(fields), schema()->metadata());
        return RecordBatch::create(std::move(new_schema), num_rows(), std::move(columns),
                                   GetSyncEvent());
    }

    std::shared_ptr<RecordBatch> RecordBatch::slice(int64_t offset) const {
        return slice(offset, this->num_rows() - offset);
    }

    std::string RecordBatch::to_string() const {
        std::stringstream ss;
        KCHECK_OK(pretty_print(*this, 0, &ss));
        return ss.str();
    }

    namespace {

        turbo::Status ValidateBatch(const RecordBatch &batch, bool full_validation) {
            for (int i = 0; i < batch.num_columns(); ++i) {
                TURBO_RETURN_NOT_OK(ValidateColumnLength(batch, i));
                const auto &array = *batch.column(i);
                const auto &schema_type = batch.schema()->field(i)->type();
                if (!array.type()->equals(schema_type)) {
                    return turbo::invalid_argument_error("Column ", i,
                                                         " type not match schema: ", array.type()->to_string(), " vs ",
                                                         schema_type->to_string());
                }
                const auto st = full_validation ? internal::ValidateArrayFull(array)
                                                : internal::ValidateArray(array);
                if (!st.ok()) {
                    return turbo::invalid_argument_error("In column ", i, ": ", st.to_string());
                }
            }
            return turbo::OkStatus();
        }

    }  // namespace

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::copy_to(
            const std::shared_ptr<MemoryManager> &to) const {
        ArrayVector copied_columns;
        copied_columns.reserve(num_columns());
        for (const auto &col: columns()) {
            TURBO_MOVE_OR_RAISE(auto c, col->copy_to(to));
            copied_columns.push_back(std::move(c));
        }

        return create(schema_, num_rows(), std::move(copied_columns));
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatch::view_or_copy_to(
            const std::shared_ptr<MemoryManager> &to) const {
        ArrayVector copied_columns;
        copied_columns.reserve(num_columns());
        for (const auto &col: columns()) {
            TURBO_MOVE_OR_RAISE(auto c, col->view_or_copy_to(to));
            copied_columns.push_back(std::move(c));
        }

        return create(schema_, num_rows(), std::move(copied_columns));
    }

    turbo::Status RecordBatch::validate() const {
        return ValidateBatch(*this, /*full_validation=*/false);
    }

    turbo::Status RecordBatch::validate_full() const {
        return ValidateBatch(*this, /*full_validation=*/true);
    }

    // ----------------------------------------------------------------------
    // Base record batch reader

    turbo::Result<RecordBatchVector> RecordBatchReader::to_record_batches() {
        RecordBatchVector batches;
        while (true) {
            std::shared_ptr<RecordBatch> batch;
            TURBO_RETURN_NOT_OK(read_next(&batch));
            if (!batch) {
                break;
            }
            batches.emplace_back(std::move(batch));
        }
        return batches;
    }

    turbo::Result<std::shared_ptr<Table>> RecordBatchReader::to_table() {
        TURBO_MOVE_OR_RAISE(auto batches, to_record_batches());
        return Table::from_record_batches(schema(), std::move(batches));
    }

    class SimpleRecordBatchReader : public RecordBatchReader {
    public:
        SimpleRecordBatchReader(turbo::Iterator<std::shared_ptr<RecordBatch>> it,
                                std::shared_ptr<Schema> schema,
                                DeviceAllocationType device_type = DeviceAllocationType::kCPU)
                : schema_(std::move(schema)), it_(std::move(it)), device_type_(device_type) {}

        SimpleRecordBatchReader(std::vector<std::shared_ptr<RecordBatch>> batches,
                                std::shared_ptr<Schema> schema,
                                DeviceAllocationType device_type = DeviceAllocationType::kCPU)
                : schema_(std::move(schema)),
                  it_(turbo::make_vector_iterator(std::move(batches))),
                  device_type_(device_type) {}

        turbo::Status read_next(std::shared_ptr<RecordBatch> *batch) override {
            return it_.next().try_value(batch);
        }

        std::shared_ptr<Schema> schema() const override { return schema_; }

        DeviceAllocationType device_type() const override { return device_type_; }

    protected:
        std::shared_ptr<Schema> schema_;
        turbo::Iterator<std::shared_ptr<RecordBatch>> it_;
        DeviceAllocationType device_type_;
    };

    turbo::Result<std::shared_ptr<RecordBatchReader>> RecordBatchReader::create(
            std::vector<std::shared_ptr<RecordBatch>> batches, std::shared_ptr<Schema> schema,
            DeviceAllocationType device_type) {
        if (schema == nullptr) {
            if (batches.size() == 0 || batches[0] == nullptr) {
                return turbo::invalid_argument_error("Cannot infer schema from empty vector or nullptr");
            }

            schema = batches[0]->schema();
        }

        return std::make_shared<SimpleRecordBatchReader>(std::move(batches), std::move(schema),
                                                         device_type);
    }

    turbo::Result<std::shared_ptr<RecordBatchReader>> RecordBatchReader::make_from_iterator(
            turbo::Iterator<std::shared_ptr<RecordBatch>> batches, std::shared_ptr<Schema> schema,
            DeviceAllocationType device_type) {
        if (schema == nullptr) {
            return turbo::invalid_argument_error("Schema cannot be nullptr");
        }

        return std::make_shared<SimpleRecordBatchReader>(std::move(batches), std::move(schema),
                                                         device_type);
    }

    RecordBatchReader::~RecordBatchReader() {
        TURBO_WARN_NOT_OK(this->close(), "Implicitly called RecordBatchReader::close failed");
    }

    turbo::Result<std::shared_ptr<RecordBatch>> concatenate_record_batches(
            const RecordBatchVector &batches, MemoryPool *pool) {
        int64_t length = 0;
        size_t n = batches.size();
        if (n == 0) {
            return turbo::invalid_argument_error("Must pass at least one recordbatch");
        }
        int cols = batches[0]->num_columns();
        auto schema = batches[0]->schema();
        for (size_t i = 0; i < batches.size(); ++i) {
            if (auto sync = batches[i]->GetSyncEvent()) {
                TURBO_RETURN_NOT_OK(sync->Wait());
            }

            length += batches[i]->num_rows();
            if (!schema->equals(batches[i]->schema())) {
                return turbo::invalid_argument_error(
                        "Schema of RecordBatch index ", i, " is ", batches[i]->schema()->to_string(),
                        ", which does not match index 0 recordbatch schema: ", schema->to_string());
            }
        }

        std::vector<std::shared_ptr<Array>> concatenated_columns;
        concatenated_columns.reserve(cols);
        for (int col = 0; col < cols; ++col) {
            ArrayVector column_arrays;
            column_arrays.reserve(batches.size());
            for (const auto &batch: batches) {
                column_arrays.emplace_back(batch->column(col));
            }
            TURBO_MOVE_OR_RAISE(auto concatenated_column, concatenate(column_arrays, pool));
            concatenated_columns.emplace_back(std::move(concatenated_column));
        }
        return RecordBatch::create(std::move(schema), length, std::move(concatenated_columns));
    }
}  // namespace nebula
